from Crypto.Cipher import AES
from Crypto.Util.Padding import pad, unpad
from Crypto.Protocol.KDF import PBKDF2
import hashlib

class AESHelper:
    """
    AES 加密/解密工具类（支持自定义 IV）
    支持 CBC/ECB/CTR 模式，PKCS7/ZeroPadding 填充，自动调整密钥长度
    """
    def __init__(self, key: bytes, key_length: int = 256, mode: str = "CBC", padding: str = "PKCS7", iv: bytes = None):
        """
        :param key: 原始密钥（任意长度）
        :param key_length: 密钥长度（128/192/256）
        :param mode: 加密模式（CBC/ECB/CTR）
        :param padding: 填充模式（PKCS7/ZeroPadding/None）
        :param iv: 初始化向量（CBC/CTR 模式必须提供，16字节）
        """
        self.key = self._adjust_key(key, key_length)
        self.key_length = key_length
        self.mode = mode.upper()
        self.padding = padding.upper()
        self.iv = iv  # 由外部传入
    def __enter__(self):
        """上下文管理器入口"""
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        """上下文管理器退出，清理敏感数据"""
        # 清零敏感数据
        if hasattr(self, 'key') and self.key:
            self.key = b''
        if hasattr(self, 'iv') and self.iv:
            self.iv = b''
        return False  # 不抑制异常
    def _adjust_key(self, key: bytes, key_length: int) -> bytes:
        """调整密钥长度（PBKDF2 哈希派生）"""
        if len(key) * 8 == key_length:
            return key
        return PBKDF2(key, b'salt', dkLen=key_length // 8, count=1000, prf=lambda p, s: hashlib.sha256(p + s).digest())

    def _get_cipher(self):
        """根据模式返回 AES 加密器"""
        if self.mode == "CBC":
            if not self.iv:
                raise ValueError("IV must be provided for CBC mode.")
            return AES.new(self.key, AES.MODE_CBC, self.iv)
        elif self.mode == "ECB":
            return AES.new(self.key, AES.MODE_ECB)
        elif self.mode == "CTR":
            if not self.iv:
                raise ValueError("Nonce must be provided for CTR mode.")
            return AES.new(self.key, AES.MODE_CTR, nonce=self.iv)
        else:
            raise ValueError("Unsupported AES mode. Use CBC/ECB/CTR.")

    def _apply_padding(self, data: bytes) -> bytes:
        """应用填充"""
        if self.padding == "PKCS7":
            return pad(data, AES.block_size)
        elif self.padding == "ZEROPADDING":
            pad_len = AES.block_size - (len(data) % AES.block_size)
            return data + bytes([0] * pad_len)
        elif self.padding == "NONE":
            if len(data) % AES.block_size != 0:
                raise ValueError("Data length must be multiple of 16 if no padding.")
            return data
        else:
            raise ValueError("Unsupported padding. Use PKCS7/ZeroPadding/None.")

    def _remove_padding(self, data: bytes) -> bytes:
        """移除填充"""
        if self.padding == "PKCS7":
            return unpad(data, AES.block_size)
        elif self.padding == "ZEROPADDING":
            return data.rstrip(b'\x00')
        elif self.padding == "NONE":
            return data
        else:
            raise ValueError("Unsupported padding. Use PKCS7/ZeroPadding/None.")

    def encrypt(self, plaintext: bytes) -> bytes:
        """加密数据"""
        cipher = self._get_cipher()
        padded_data = self._apply_padding(plaintext)
        encrypted = cipher.encrypt(padded_data)
        return encrypted  # 不自动拼接 IV，由调用者管理

    def decrypt(self, ciphertext: bytes) -> bytes:
        """解密数据"""
        cipher = self._get_cipher()
        decrypted = cipher.decrypt(ciphertext)
        return self._remove_padding(decrypted)




